import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from nav_msgs.msg import Odometry
from geometry_msgs.msg import PoseWithCovarianceStamped
import numpy as np
from tf_transformations import euler_from_quaternion, quaternion_from_euler
from yesense_interface.msg import ImuData # type: ignore

class ExtendedKalmanFilter:
    """
    扩展卡尔曼滤波器实现，用于融合IMU和轮式里程计数据
    
    状态向量: [x, y, theta, v, omega]
    x, y: 位置坐标
    theta: 航向角
    v: 线速度
    omega: 角速度
    """
    def __init__(self, dt):
        # 初始化状态估计和协方差矩阵
        self.x = np.zeros(5)  # 状态向量 [x, y, theta, v, omega]
        self.P = np.eye(5)    # 状态协方差矩阵，初始为单位矩阵
        
        # 过程噪声协方差矩阵 - 调整这些参数可影响预测的信任度
        self.Q = np.diag([0.01, 0.01, 0.01, 0.1, 0.1])
        
        # 测量噪声协方差矩阵 - 调整这些参数可影响传感器测量的信任度
        self.R_imu = np.diag([0.1, 0.1, 0.01])  # IMU测量噪声
        self.R_odom = np.diag([0.1, 0.1, 0.01]) # 轮式里程计测量噪声
        
        self.dt = dt  # 时间步长，用于状态预测
    
    def predict(self, u):
        """
        EKF预测步骤：基于系统模型预测下一时刻的状态
        
        参数:
            u: 控制输入向量 [加速度, 角加速度]
        """
        x, y, theta, v, omega = self.x
        
        # 非线性状态转移方程 - 基于当前状态和控制输入预测下一状态
        theta_new = theta + omega * self.dt
        x_new = x + v * np.cos(theta_new) * self.dt  # 考虑航向角的位置更新
        y_new = y + v * np.sin(theta_new) * self.dt
        v_new = v + u[0] * self.dt  # 线速度更新（基于加速度）
        omega_new = omega + u[1] * self.dt  # 角速度更新（基于角加速度）
        
        # 更新状态估计
        self.x = np.array([x_new, y_new, theta_new, v_new, omega_new])
        
        # 计算状态转移雅可比矩阵（状态方程的一阶泰勒展开）
        F = np.eye(5)
        # 位置对航向角和速度的偏导数
        F[0, 2] = -v * np.sin(theta_new) * self.dt
        F[0, 3] = np.cos(theta_new) * self.dt
        F[1, 2] = v * np.cos(theta_new) * self.dt
        F[1, 3] = np.sin(theta_new) * self.dt
        
        # 更新状态协方差矩阵
        self.P = F @ self.P @ F.T + self.Q
        
        return self.x
    
    def update_imu(self, z):
        """
        使用IMU数据更新EKF状态估计
        
        参数:
            z: IMU测量向量 [ax, ay, yaw_rate]
        """
        # 测量矩阵 - 将状态向量映射到IMU测量空间
        H = np.zeros((3, 5))
        H[0, 3] = 1  # ax与线速度v相关
        H[1, 4] = 1  # ay与角速度omega相关
        H[2, 2] = 1  # yaw_rate与航向角theta相关
        
        # 卡尔曼增益计算 - 权衡预测和测量的可信度
        S = H @ self.P @ H.T + self.R_imu  # 测量残差协方差
        K = self.P @ H.T @ np.linalg.inv(S)  # 卡尔曼增益
        
        # 状态更新 - 结合预测和测量
        y = z - np.array([self.x[3], self.x[4], self.x[2]])  # 测量残差
        self.x = self.x + K @ y  # 状态更新
        
        # 协方差更新
        I = np.eye(5)
        self.P = (I - K @ H) @ self.P
        
        return self.x
    
    def update_odom(self, z):
        """
        使用轮式里程计数据更新EKF状态估计
        
        参数:
            z: 轮式里程计测量向量 [v, omega, yaw]
        """
        # 测量矩阵 - 将状态向量映射到里程计测量空间
        H = np.zeros((3, 5))
        H[0, 3] = 1  # 线速度测量
        H[1, 4] = 1  # 角速度测量
        H[2, 2] = 1  # 航向角测量
        
        # 卡尔曼增益计算
        S = H @ self.P @ H.T + self.R_odom  # 测量残差协方差
        K = self.P @ H.T @ np.linalg.inv(S)  # 卡尔曼增益
        
        # 状态更新
        y = z - np.array([self.x[3], self.x[4], self.x[2]])  # 测量残差
        self.x = self.x + K @ y  # 状态更新
        
        # 协方差更新
        I = np.eye(5)
        self.P = (I - K @ H) @ self.P
        
        return self.x

class EKFFusionNode(Node):
    """
    ROS2节点：使用扩展卡尔曼滤波器融合IMU和轮式里程计数据
    """
    def __init__(self):
        super().__init__('ekf_imu_odom_fusion_node')
        
        # 声明ROS2参数，允许通过命令行或配置文件修改
        self.declare_parameter('imu_topic', '/imu_data')
        self.declare_parameter('odom_topic', '/odom')
        self.declare_parameter('output_topic', '/ekf_odom')
        self.declare_parameter('dt', 0.01)
        
        # 获取参数值
        imu_topic = self.get_parameter('imu_topic').get_parameter_value().string_value
        odom_topic = self.get_parameter('odom_topic').get_parameter_value().string_value
        output_topic = self.get_parameter('output_topic').get_parameter_value().string_value
        dt = self.get_parameter('dt').get_parameter_value().double_value
        
        # 创建EKF实例
        self.ekf = ExtendedKalmanFilter(dt)
        
        # 创建ROS2订阅者
        self.imu_sub = self.create_subscription(
            ImuData,
            imu_topic,
            self.imu_callback,
            10)
        
        self.odom_sub = self.create_subscription(
            Odometry,
            odom_topic,
            self.odom_callback,
            10)
        
        # 创建ROS2发布者 - 发布融合后的里程计数据
        self.ekf_odom_pub = self.create_publisher(
            Odometry,
            output_topic,
            10)
        
        # 存储最新的传感器数据
        self.last_imu_msg = None
        self.last_odom_msg = None
        
        # 创建定时器，按固定频率执行融合算法
        self.timer = self.create_timer(dt, self.timer_callback)
        
        self.get_logger().info('EKF融合节点已启动')
    
    def imu_callback(self, msg):
        """IMU消息回调函数"""
        self.last_imu_msg = msg
    
    def odom_callback(self, msg):
        """里程计消息回调函数"""
        self.last_odom_msg = msg
    
    def timer_callback(self):
        """定时器回调函数 - 执行EKF融合算法"""
        #检查是否接收到了所有必要的传感器数据
        if self.last_imu_msg is None or self.last_odom_msg is None:
            return
        print("开始")
        # 从IMU消息中提取数据
        imu_data = self.last_imu_msg
        ax = imu_data.acc.x  # 线性加速度x分量
        ay = imu_data.acc.y  # 线性加速度y分量
        yaw_rate = imu_data.acc.z  # 偏航角速度
        
        # 从里程计消息中提取数据
        odom_data = self.last_odom_msg
        # 计算线速度大小
        v = np.sqrt(odom_data.twist.twist.linear.x**2 + 
                   odom_data.twist.twist.linear.y**2)
        
        # 从四元数中提取偏航角
        orientation = odom_data.pose.pose.orientation
        quat = [orientation.x, orientation.y, orientation.z, orientation.w]
        _, _, yaw = euler_from_quaternion(quat)
        
        # 构建EKF输入和测量向量
        u = np.array([ax, yaw_rate])  # 控制输入向量 [加速度, 角加速度]
        z_imu = np.array([ax, ay, yaw_rate])  # IMU测量向量
        z_odom = np.array([v, odom_data.twist.twist.angular.z, yaw])  # 里程计测量向量
        
        # EKF预测和更新步骤
        self.ekf.predict(u)  # 基于控制输入预测下一状态
        self.ekf.update_imu(z_imu)  # 用IMU数据更新状态估计
        self.ekf.update_odom(z_odom)  # 用里程计数据更新状态估计
        
        # 发布融合后的里程计消息
        ekf_odom = Odometry()
        ekf_odom.header = odom_data.header  # 使用原始里程计的时间戳和坐标系
        ekf_odom.child_frame_id = odom_data.child_frame_id
        
        # 设置位置信息
        ekf_odom.pose.pose.position.x = self.ekf.x[0]  # x坐标
        ekf_odom.pose.pose.position.y = self.ekf.x[1]  # y坐标
        
        # 将欧拉角转换为四元数表示朝向
        quat = quaternion_from_euler(0, 0, self.ekf.x[2])  # 从偏航角创建四元数
        ekf_odom.pose.pose.orientation.x = quat[0]
        ekf_odom.pose.pose.orientation.y = quat[1]
        ekf_odom.pose.pose.orientation.z = quat[2]
        ekf_odom.pose.pose.orientation.w = quat[3]
        
        # 设置速度信息
        ekf_odom.twist.twist.linear.x = self.ekf.x[3]  # 线速度
        ekf_odom.twist.twist.angular.z = self.ekf.x[4]  # 角速度
        print("处理好了")
        # 发布融合后的里程计消息
        self.ekf_odom_pub.publish(ekf_odom)

def main(args=None):
    """ROS2节点主函数"""
    rclpy.init(args=args)
    node = EKFFusionNode()
    rclpy.spin(node)  # 进入节点循环，等待回调
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()